#include "os.h"

extern void switch_to(struct context *next);

#define STACK_SIZE  1024
#define MAX_TASKS   10
#define MAX_PRIORITY 256

struct task {
    struct context ctx;      	// 任务上下文
    uint8_t priority;        	// 任务优先级
    uint32_t timeslice;      	// 时间片大小
    uint32_t remaining_ticks; 	// 剩余时间片
    int active;              	// 是否活跃（1: 活跃, 0: 已退出）
};

uint8_t __attribute__((aligned(16))) task_stack[MAX_TASKS][STACK_SIZE];
struct task tasks[MAX_TASKS];

static int _top = 0;
static int _current = -1;

void sched_init()
{
    w_mscratch(0);
    w_mie(r_mie() | MIE_MSIE);
}

/*
 * DESCRIPTION
 *  Create a task.
 *  - task: 任务入口函数
 *  - param: 传递给任务的参数
 *  - priority: 任务优先级
 *  - timeslice: 时间片大小
 * RETURN VALUE
 *  0: success
 * -1: if error occurred
 */
int task_create(void (*task)(void* param), void *param, uint8_t priority, uint32_t timeslice)
{
    if (_top < MAX_TASKS) {
        tasks[_top].ctx.sp = (reg_t) &task_stack[_top][STACK_SIZE];
        tasks[_top].ctx.pc = (reg_t) task;
        tasks[_top].ctx.a0 = (reg_t) param; 		// 将参数传递给任务
        tasks[_top].priority = priority;
        tasks[_top].timeslice = timeslice;
        tasks[_top].remaining_ticks = timeslice; 	// 初始化剩余时间片
        tasks[_top].active = 1; 					// 标记任务为活跃
        _top++;
        return 0;
    } else {
        return -1;
    }
}


/*
 * DESCRIPTION
 *  Implement a priority-based scheduler with round-robin for tasks of the same priority.
 *  Add timeslice management.
 */
void schedule(void)
{
    if (_top <= 0) {
        panic("Num of task should be greater than zero!");
        return;
    }

    int highest_priority = MAX_PRIORITY;
    int next_task = -1;

    // 动态优先级调整：逐渐提升未运行任务的优先级
    for (int i = 0; i < _top; i++) {
        if (tasks[i].active && i != _current) {
            tasks[i].priority = (tasks[i].priority > 0) ? tasks[i].priority - 1 : 0;
        }
    }

    // 遍历任务列表，选择优先级最高的任务
    for (int i = 0; i < _top; i++) {
        if (tasks[i].active) {
            if (next_task == -1 || tasks[i].priority < highest_priority || 
                (tasks[i].priority == highest_priority && i > _current)) {
                highest_priority = tasks[i].priority;
                next_task = i;
            }
        }
    }

    if (next_task == -1) {
        panic("No active tasks to schedule!");
        return;
    }

    // 时间片管理
    if (_current >= 0 && _current < _top) {
        tasks[_current].remaining_ticks--;
        if (tasks[_current].remaining_ticks == 0) {
			// 重置时间片
            tasks[_current].remaining_ticks = tasks[_current].timeslice;
        } else if (tasks[_current].priority <= tasks[next_task].priority) {
            // 当前任务时间片未用完，且优先级不低于下一个任务，继续运行
            return;
        }
    }

    // 切换到下一个任务
    _current = next_task;
    struct context *next = &(tasks[_current].ctx);
    switch_to(next);
}

/*
 * DESCRIPTION
 *  Exit the current task.
 */
void task_exit(void)
{
    if (_current >= 0 && _current < _top) {
 		/* no activate */
        tasks[_current].active = 0; 
    	/* trigger a machine-level software interrupt */
		int id = r_mhartid();
		*(uint32_t*)CLINT_MSIP(id) = 1;
    }
}

/*
 * DESCRIPTION
 *  task_yield() causes the calling task to relinquish the CPU and a new 
 *  task gets to run.
 */
void task_yield()
{
    if (_current >= 0 && _current < _top) {
        tasks[_current].remaining_ticks = 0; 
    }
    /* trigger a machine-level software interrupt */
    int id = r_mhartid();
    *(uint32_t*)CLINT_MSIP(id) = 1;
}

/*
 * a very rough implementation, just to consume the cpu
 */
void task_delay(volatile int count)
{
    count *= 50000;
    while (count--);
}
